package main

import (
	"fmt"
	"net"
	"os"
	"strconv"

	. "gitee.com/shuohe/keytool/msg"

	"gitee.com/shuohe/xutils/xcrypto"
	"gitee.com/shuohe/xutils/xnet"
	"gitee.com/shuohe/xutils/xlog"
	"gitee.com/shuohe/xutils/xbuf"
	"gitee.com/shuohe/xutils/xerror"
)

var g_keysetPath = ""
var g_keyset     = xcrypto.NewKeySet()
var g_mkey       xcrypto.ISecretKey
var g_maxkidx    = 0

//=============================================================================

type SvChCallback struct {
	Channel  *xnet.TCPChannel
	xb1      *xbuf.Buffer		// For temporary use
	xb2      *xbuf.Buffer		// For temporary use
	xb3      *xbuf.Buffer		// ONLY used by SvChCallback.writeMessage()
}

func (this *SvChCallback) writeMessage(msg *Bitmap32Message) error {
	this.xb3.Clear()
	msg.MarshalX(this.xb3)
	return this.Channel.Write(this.xb3.Bytes())
}

func (this *SvChCallback) encryptKey(key xcrypto.IKey) []byte {
	this.xb1.Consume()
	key.ExportX(this.xb1)

	this.xb2.Consume()
	g_mkey.EncryptX(this.xb1.Bytes(), this.xb2)

	return this.xb2.Bytes()
}

func (this *SvChCallback) exportKey(key xcrypto.IKey) []byte {
    this.xb1.Consume()
    key.ExportX(this.xb1)

    return this.xb1.Bytes()
}

func (this *SvChCallback) SetChannel(ch *xnet.TCPChannel) {
	this.Channel = ch
}

func (this *SvChCallback) OnConnect() error {
	xlog.Info(this.Channel.FullName(), "is connected.")
	this.xb1 = xbuf.New(1024)
	this.xb2 = xbuf.New(1024)
	this.xb3 = xbuf.New(1024)

	go this.Channel.Serve()

	return nil
}

func (this *SvChCallback) OnMessage(msg []byte) error {
	xlog.Info(this.Channel.FullName(), "message arrived.")

	req, err := Bitmap32Message_Unmarshal(msg)
	if err != nil {
		xlog.Errorfe(err, "Corrupted message data.")
		return xerror.Wrap(err, "Corrupted message data.")
	}

	xlog.Info("Request message:")
	req.Dump(xlog.DefaultLogger.Writer(), "--> ")

	this.xb1.Clear()
	this.xb2.Clear()

	var rsp *Bitmap32Message

	switch req.GetString(MsgType) {
	case MsgType_GenkeyRequest:
		rsp, err = this.genkey(req)
	case MsgType_GetkeyRequest:
		rsp, err = this.getkey(req)
	case MsgType_EncryptRequest:
		rsp, err = this.encrypt(req)
	case MsgType_SignRequest:
		rsp, err = this.sign(req)
	case MsgType_VerifyRequest:
		rsp, err = this.verify(req)
	default:
		rsp = &Bitmap32Message{}
		rsp.SetString(MsgType, req.GetString(MsgType))
		rsp.SetString(RetCode, RetCode_EMsgType)
		err = nil
	}
	if err != nil {
		// Will close this channel
		return xerror.Wrap(err, "Internal error.")
	}

	xlog.Info("Response message:")
	rsp.Dump(xlog.DefaultLogger.Writer(), "<-- ")

	return this.writeMessage(rsp)
}

func (this *SvChCallback) OnClose() {
	xlog.Info(this.Channel.FullName(), "is closed.")
}

func (this *SvChCallback) OnError(err error) {
	xlog.Errorfe(err, this.Channel.FullName() + " has error.")
}

func (this *SvChCallback) genkey(req *Bitmap32Message) (*Bitmap32Message, error) {
	rsp := &Bitmap32Message{}
	rsp.SetString(MsgType, MsgType_GenkeyResponse)

	// err := req.ValidateFields()
	// TODO: Validate message format
	
	options := req.GetString(Options)
	optSave := options[0] != '0'
	optEcho := !optSave || options[1] != '0'

	k, err := xcrypto.GenerateKey(req.GetString(KeyAlgo))
	if err != nil {
		xlog.Errorfe(err, "Failed to generate key.")
		rsp.SetString(RetCode, RetCode_EKeyAlgo)
		return rsp, nil
	}

	if optSave {
		g_maxkidx++
		g_keyset.Add(strconv.Itoa(g_maxkidx), k)
		g_keyset.Save(g_keysetPath)
		rsp.SetInt(KeyIndex1, int64(g_maxkidx))
	}

	if k.Type() == xcrypto.KeyType_Private {
		rsp.SetBytes(PublicKey, this.exportKey(k.(xcrypto.IPrivateKey).Public()))
	}

	if optEcho {
		kdata := this.encryptKey(k)
		if k.Type() == xcrypto.KeyType_Secret {
			rsp.SetBytes(SecretKey1, kdata)
		} else {
			rsp.SetBytes(PrivateKey, kdata)
		}
	}
	rsp.SetString(RetCode, RetCode_Ok)

	return rsp, nil
}

func (this *SvChCallback) getkey(req *Bitmap32Message) (*Bitmap32Message, error) {
	rsp := &Bitmap32Message{}
	rsp.SetString(MsgType, MsgType_GetkeyResponse)

	// err := req.ValidateFields()
	// TODO: Validate message format

	item, err := g_keyset.Get(strconv.Itoa(int(req.GetInt(KeyIndex1))))
	if err != nil {
		xlog.Errorfe(err, "Failed to get key from the keyset.")
		rsp.SetString(RetCode, RetCode_EKeyIndex1)
		return rsp, nil
	}

	switch req.GetString(Options)[0] {
	case '1': // Secret
		if item.Key.Type() != xcrypto.KeyType_Secret {
			xlog.Error("Incorrect key type.")
			rsp.SetString(RetCode, RetCode_EKeyType)
			return rsp, nil
		}
		rsp.SetBytes(SecretKey1, this.encryptKey(item.Key))
	case '2': // Private
		if item.Key.Type() != xcrypto.KeyType_Private {
			xlog.Error("Incorrect key type.")
			rsp.SetString(RetCode, RetCode_EKeyType)
			return rsp, nil
		}
		rsp.SetBytes(PrivateKey, this.encryptKey(item.Key))
        rsp.SetBytes(PublicKey, this.exportKey(item.Key.(xcrypto.IPrivateKey).Public()))
	case '3':  // Public
		if item.Key.Type() != xcrypto.KeyType_Private {
			xlog.Error("Incorrect key type.")
			rsp.SetString(RetCode, RetCode_EKeyType)
			return rsp, nil
		}
		rsp.SetBytes(PublicKey, this.exportKey(item.Key))
	default: // Auto
		if item.Key.Type() == xcrypto.KeyType_Secret {
			rsp.SetBytes(SecretKey1, this.encryptKey(item.Key))
		} else { // Private
			rsp.SetBytes(PrivateKey, this.encryptKey(item.Key))
			rsp.SetBytes(PublicKey, this.exportKey(item.Key.(xcrypto.IPrivateKey).Public()))
		}
	}
	rsp.SetString(RetCode, RetCode_Ok)

	return rsp, nil
}

func (this *SvChCallback) encrypt(req *Bitmap32Message) (*Bitmap32Message, error) {
	rsp := &Bitmap32Message{}
	rsp.SetString(MsgType, MsgType_EncryptResponse)

	// err := req.ValidateFields()
	// TODO: Validate message format

	options := req.GetString(Options)
	optUseIndex := options[0] == '0'
	optEnc := options[1] == '0'

	var sk xcrypto.IKey

	if optUseIndex {
		item, err := g_keyset.Get(strconv.Itoa(int(req.GetInt(KeyIndex1))))
		if err != nil {
			xlog.Errorfe(err, "Invalid key index.")
			rsp.SetString(RetCode, RetCode_EKeyIndex1)
			return rsp, nil
		}

		sk = item.Key
	} else {
		this.xb1.Clear()
		err := g_mkey.DecryptX(req.GetBytes(SecretKey1), this.xb1)
		if err != nil {
			xlog.Errorfe(err, "Can not recover secret key.")
			rsp.SetString(RetCode, RetCode_ESecretKey1)
			return rsp, nil
		}

		sk, err = xcrypto.LoadKey(this.xb1.Bytes())
		if err != nil {
			xlog.Errorfe(err, "Can not create secret key object.")
			rsp.SetString(RetCode, RetCode_ESecretKey1)
			return rsp, nil
		}
	}

	if sk.Type() != xcrypto.KeyType_Secret {
		xlog.Error("Incorrect key type.")
		rsp.SetString(RetCode, RetCode_EKeyType)
		return rsp, nil
	}

	this.xb1.Clear()
	if optEnc {
		sk.(xcrypto.ISecretKey).EncryptX(req.GetBytes(Meterial), this.xb1)
	} else {
		err := sk.(xcrypto.ISecretKey).DecryptX(req.GetBytes(Meterial), this.xb1)
		if err != nil {
			xlog.Error("Can not decrypt invalid meterial.")
			rsp.SetString(RetCode, RetCode_EMeterial)
			return rsp, nil
		}
	}

	rsp.SetBytes(Meterial, this.xb1.Bytes())
	rsp.SetString(RetCode, RetCode_Ok)

	return rsp, nil
}

func (this *SvChCallback) sign(req *Bitmap32Message) (*Bitmap32Message, error) {
	rsp := &Bitmap32Message{}
	rsp.SetString(MsgType, MsgType_SignResponse)

	// err := req.ValidateFields()
	// TODO: Validate message format

	var sk xcrypto.IKey

	optUseIndex := req.GetString(Options)[0] == '0'
	if optUseIndex {
		item, err := g_keyset.Get(strconv.Itoa(int(req.GetInt(KeyIndex1))))
		if err != nil {
			xlog.Errorfe(err, "Invalid key index.")
			rsp.SetString(RetCode, RetCode_EKeyIndex1)
			return rsp, nil
		}

		sk = item.Key
	} else {
		this.xb1.Clear()
		err := g_mkey.DecryptX(req.GetBytes(PrivateKey), this.xb1)
		if err != nil {
			xlog.Errorfe(err, "Can not recover private key.")
			rsp.SetString(RetCode, RetCode_EPrivateKey)
			return rsp, nil
		}

		sk, err = xcrypto.LoadKey(this.xb1.Bytes())
		if err != nil {
			xlog.Errorfe(err, "Can not create private key object.")
			rsp.SetString(RetCode, RetCode_EPrivateKey)
			return rsp, nil
		}
	}

	if sk.Type() != xcrypto.KeyType_Private {
		xlog.Error("Incorrect key type.")
		rsp.SetString(RetCode, RetCode_EKeyType)
		return rsp, nil
	}

	s := sk.(xcrypto.IPrivateKey).Sign(req.GetBytes(Meterial))

	rsp.SetBytes(Signature, s)
	rsp.SetString(RetCode, RetCode_Ok)

	return rsp, nil
}

func (this *SvChCallback) verify(req *Bitmap32Message) (*Bitmap32Message, error) {
	rsp := &Bitmap32Message{}
	rsp.SetString(MsgType, MsgType_VerifyResponse)

	// err := req.ValidateFields()
	// TODO: Validate message format

	var pk xcrypto.IKey
	var err error

	optUseIndex := req.GetString(Options)[0] == '0'
	if optUseIndex {
		item, err := g_keyset.Get(strconv.Itoa(int(req.GetInt(KeyIndex1))))
		if err != nil {
			xlog.Errorfe(err, "Invalid key index.")
			rsp.SetString(RetCode, RetCode_EKeyIndex1)
			return rsp, nil
		}

		pk = item.Key
	} else {
		pk, err = xcrypto.LoadKey(req.GetBytes(PublicKey))
		if err != nil {
			xlog.Errorfe(err, "Can not create public key object.")
			rsp.SetString(RetCode, RetCode_EPublicKey)
			return rsp, nil
		}
	}

	if pk.Type() == xcrypto.KeyType_Private {
		pk = pk.(xcrypto.IPrivateKey).Public()
	} else if pk.Type() != xcrypto.KeyType_Public {
		xlog.Error("Incorrect key type.")
		rsp.SetString(RetCode, RetCode_EKeyType)
		return rsp, nil
	}

	ok := pk.(xcrypto.IPublicKey).Verify(req.GetBytes(Meterial), req.GetBytes(Signature))
	if ok {
		rsp.SetString(VerifyOk, "1")
	} else {
		rsp.SetString(VerifyOk, "0")
	}
	rsp.SetString(RetCode, RetCode_Ok)

	return rsp, nil
}

func NewChannel(conn *net.TCPConn, svr *xnet.TCPServer, id int) (*xnet.TCPChannel, error) {
	ch := xnet.NewTCPChannel(conn, &SvChCallback{})
	ch.Server = svr
	ch.Id = id
	//ch.IdleTimeout = 2 * time.Second

	return ch, nil
}

//=============================================================================

type SvCallback struct {
	Server		*xnet.TCPServer
}

func (this *SvCallback) SetServer(s *xnet.TCPServer) {
	this.Server = s
}

func (this *SvCallback) OnListen() {
	xlog.Info("Listening at", this.Server.ListenAddr, "...")
}

func (this *SvCallback) OnAccept(ch *xnet.TCPChannel) error {
	xlog.Info(ch.FullName(), "from", ch.RemoteAddr(), "is accepted.")
	return nil
}

func (this *SvCallback) OnClose() {
	xlog.Info("TCPServer is closed.")
}

//=============================================================================

func loadKeyset(kspath string) {
    if kspath == "" {
        kspath = "./keytool.ks"
        fmt.Printf("Hint: Default keyset path '%s' is used.\n", kspath)
    }
    n, err := g_keyset.Load(kspath)
    if err != nil {
        printErr(err, "Can not load keyset file.")
        os.Exit(1)
    }
    g_keysetPath = kspath
    if n == 0 {
        // 如果是空 keyset 文件，自动生成一个主密钥。
        k, _ := xcrypto.GenerateKey("AES-128")
        g_mkey = k.(xcrypto.ISecretKey)
        g_keyset.Add("MASTER_KEY", g_mkey)
        g_keyset.Save(g_keysetPath)
    } else {
        // 非空 keyset 文件：读取主密钥。
        k, err := g_keyset.Get("MASTER_KEY")
        if err != nil {
            printErr(err, "Can not load master key. The keyset file is not for keytool program.")
            os.Exit(1)
        }
        if k.Key.Type() != xcrypto.KeyType_Secret {
            fmt.Println("Error: Invalid master key. The keyset file is not for keytool program.")
            os.Exit(1)
        }
        g_mkey = k.Key.(xcrypto.ISecretKey)
        // 取得最大的密钥索引值
        g_keyset.Iterate(func (ki *xcrypto.KeyItem) bool {
            i, err := strconv.Atoi(ki.Alias)
            if err == nil && i > g_maxkidx {
                g_maxkidx = i
            }
            return true
        })
    }
}

func RunAsServer(conf *ServerConfig) {
	/*
	 * 加载 Keyset 文件。
	 */
	loadKeyset(conf.KeySetPath)

	/*
	 * 启动网络。
	 */
	if conf.ListenAddr == "" {
		conf.ListenAddr = ":2345"
		fmt.Printf("Hint: Default listen address '%s' is used.\n", conf.ListenAddr)
	}
	laddr, err := net.ResolveTCPAddr("tcp4", conf.ListenAddr)
	if err != nil {
		printErr(err, "Invalid listen adddress.")
		os.Exit(1)
	}
	svr := xnet.NewTCPServer(laddr, &SvCallback{}, NewChannel)
	svr.Serve()
}
